import os
import ssl
import certifi
import json
from typing import Generator, Union, Dict, Any
from io import BytesIO
from phi.agent import Agent
from phi.knowledge.pdf import PDFKnowledgeBase
from phi.knowledge.docx import DocxKnowledgeBase
from phi.knowledge.combined import CombinedKnowledgeBase
from phi.tools.duckduckgo import DuckDuckGo
from phi.vectordb.pgvector import PgVector2
from phi.model.openai import OpenAIChat
from neo4j import GraphDatabase
import httpx
import openai
from docx import Document
import PyPDF2


# Neo4j 配置
NEO4J_URI = "bolt://localhost:7687"
NEO4J_USER = "neo4j"
NEO4J_PASSWORD = "password"

# 向量数据库配置
try:
    print("正在连接向量数据库...")
    vector_db = PgVector2(
        collection="knowledge_base",
        db_url="postgresql+psycopg2://ai:ai@localhost:5432/ai"
    )
    # 测试连接
    if not vector_db.exists():
        vector_db.create()    # 如果表不存在则创建
        print("向量数据库表创建成功!")
    else:
        print("向量数据库表已存在!")
    print("向量数据库连接成功!")
except Exception as e:
    print(f"向量数据库连接失败: {str(e)}")
    print("请确保：")
    print("1. PostgreSQL (pgvector) 数据库已启动")
    print("2. 数据库连接信息正确")
    print("3. 已执行 docker run 命令启动数据库")
    raise

# 初始化知识库
knowledge_base = DocxKnowledgeBase(
    path="data/docs",
    vector_db=vector_db,
    chunk_size=1000,
    chunk_overlap=200
)
# 创建 RAG Agent
knowledge_agent = Agent(
    llm=OpenAIChat(
        model="gpt-4",
        temperature=0.7,
        max_tokens=10000
    ),
    # 添加知识库到 agent
    knowledge=knowledge_base,
    show_tool_calls=True,
    markdown=True,
    description="You are a knowledge expert that helps answer questions about documents.",
    instructions=[
        "You are an expert in answering questions about documents.",
        "1. Use the knowledge base to find relevant information",
        "2. Provide detailed answers with references to the source documents",
        "3. If information is not available, clearly state that",
        "4. Keep the answers concise but informative"
    ]
)

# 创建 Neo4j 图数据库连接
class Neo4jKnowledgeGraph:
    def __init__(self, uri, username, password):
        self.driver = GraphDatabase.driver(uri, auth=(username, password))
        # 测试连接
        with self.driver.session() as session:
            session.run("RETURN 1")
        print("Neo4j连接成功!")
    
    def close(self):
        self.driver.close()

graph_db = Neo4jKnowledgeGraph(
    uri=NEO4J_URI,
    username=NEO4J_USER,
    password=NEO4J_PASSWORD
)

def query_knowledge(question: str) -> str:
    """使用 RAG Agent 查询知识"""
    try:
        print("\n=== 开始查询知识 ===")
        
        # 使用 RAG Agent 进行查询
        response = knowledge_agent.run(f"""
        请回答以下问题。如果知识库中没有相关信息，请明确说明。
        如果有相关信息，请提供详细的回答，并引用具体的文档内容。

        问题：{question}
        """)

        # 处理响应
        if isinstance(response, str):
            return response
        elif hasattr(response, 'content'):
            return response.content
        else:
            return str(response)
        
    except Exception as e:
        print(f"查询知识时出错: {str(e)}")
        return "抱歉，查询过程中出现错误。"

def process_document(file_path: str, partition: str) -> bool:
    try:
        print(f"处理文档: {file_path}")
        
        # 加载文档到知识库
        knowledge_base.load()
        
        # 从文档中提取知识并存储到Neo4j
        with open(file_path, 'rb') as f:
            content = f.read()
            if file_path.endswith('.docx'):
                # 使用 BytesIO 处理 Word 文档
                doc = Document(BytesIO(content))
                text = '\n'.join([paragraph.text for paragraph in doc.paragraphs])
            elif file_path.endswith('.pdf'):
                # 处理 PDF 文档
                reader = PyPDF2.PdfReader(BytesIO(content))
                text = '\n'.join([page.extract_text() for page in reader.pages])
            else:
                print(f"不支持的文件类型: {file_path}")
                return False
                
        # 提取并存储知识
        success = extract_and_store_knowledge(text, partition)
        if not success:
            print("知识提取失败")
            return False
            
        print("文档处理成功")
        return True
        
    except Exception as e:
        print(f"处理文档时出错: {str(e)}")
        return False

def extract_and_store_knowledge(text: str, partition: str) -> bool:
    """从文本中提取并存储知识到Neo4j（仅用于知识图谱可视化）"""
    try:
        print("\n=== 开始提取知识用于可视化 ===")
        print(f"文本长度: {len(text)} 字符")
        # 使用 knowledge_agent 提取知识
        prompt = f"""
        请从以下文本中提取全文所有的实体和关系，并按JSON格式返回。
        只返回JSON格式的数据，不要有其他说明文字。

        实体类型必须是以下之一：
        - Product: 产品
        - Feature: 特性
        - Application: 应用场景
        - Technology: 技术
        - Material: 材料
        - Parameter: 参数

        关系类型必须是以下之一：
        - HAS_FEATURE: 产品具有的特性
        - HAS_PARAMETER: 产品的技术参数
        - USED_IN: 产品的应用场景
        - MADE_OF: 产品的原材料
        - PRODUCED_BY: 产品的生产工艺
        - MAINTAINED_BY: 产品的保养方法

        文本内容：
        {text}

        请返回如下格式的JSON：
        {{
            "entities": [
                {{
                    "name": "实体名称",
                    "type": "实体类型",
                    "description": "实体描述"
                }}
            ],
            "relations": [
                {{
                    "source": "源实体名称",
                    "target": "目标实体名称",
                    "type": "关系类型",
                    "description": "关系描述"
                }}
            ]
        }}
        """
        
        # 使用 knowledge_agent 进行知识提取
        response = knowledge_agent.run(prompt)
        
        # 处理响应
        content = ""
        if isinstance(response, str):
            content = response
        elif hasattr(response, 'content'):
            content = response.content
        else:
            content = str(response)
            
        print(f"\n获取到的响应: {content}")
        
        # 提取JSON部分
        try:
            # 尝试直接解析整个响应
            knowledge = json.loads(content)
        except json.JSONDecodeError:
            # 如果失败，尝试提取 JSON 部分
            start_idx = content.find('{')
            end_idx = content.rfind('}') + 1
            if start_idx == -1 or end_idx == 0:
                raise ValueError("响应中未找到有效的JSON数据")
            json_str = content[start_idx:end_idx]
            knowledge = json.loads(json_str)
        
        print("\n=== 解析到的知识 ===")
        print(f"实体数量: {len(knowledge.get('entities', []))}")
        print(f"关系数量: {len(knowledge.get('relations', []))}")
        
        # 验证数据结构
        if 'entities' not in knowledge or 'relations' not in knowledge:
            raise ValueError("响应数据缺少必要的字段")
        
        # 存储到Neo4j（仅用于可视化）
        with graph_db.driver.session() as session:
            # 创建实体
            print("\n=== 开始创建实体 ===")
            for entity in knowledge["entities"]:
                print(f"创建实体: {entity['name']} (类型: {entity['type']})")
                session.run(f"""
                    MERGE (n: {partition} {{name: $name}})
                    SET n.type = $type,
                        n.description = $description
                """,
                    name=entity["name"], 
                    type=entity["type"],
                    description=entity.get("description", "")
                )
            
            # 创建关系
            print("\n=== 开始创建关系 ===")
            for relation in knowledge["relations"]:
                print(f"创建关系: {relation['source']} -> {relation['target']} ({relation['type']})")
                session.run(f"""
                    MATCH (source: {partition} {{name: $source}})
                    MATCH (target: {partition} {{name: $target}})
                    OPTIONAL MATCH (source)-[r:RELATED]->(target)
                    DELETE r
                    CREATE (source)-[new_r:RELATED {{
                        type: $type,
                        description: $description
                    }}]->(target)
                """,
                    source=relation["source"],
                    target=relation["target"],
                    type=relation["type"],
                    description=relation.get("description", "")
                )
        
        print("\n=== 知识提取和存储完成 ===")
        return True
        
    except Exception as e:
        print(f"\n提取或存储知识时出错: {str(e)}")
        return False

def get_question_category(question:str):
    category = knowledge_agent.run(f"""
    问题：{question}
    请问该问题是以下哪种类型，只需回答类型即可，不带其他任何包括标点符号：
    生活，科学，军事，人文
    """)
    print('问题类型',category.content)
    return category.content

#思考用户可能会想知道的
async def think_user_ask(data):

    like_know = knowledge_agent.run(f"""
    已经知道了用户有着这些问题搜索记录：{data}
    假设你是正和用户面对面交流，你根据这些记录推测一下他想知道些什么然后回答他，
    回答他的类似句式：你可能想知道...。然后对其想知道的进行展开说说。不需要回答用户搜索记录的问题，
    你根据他搜索的猜一猜他还想知道的别的问题来进行回答。只需猜测一个问题来进行回答就行。
    """)
    print('根据用户搜索回答：',like_know.content)
    return like_know.content

def load_knowledge():
    try:
        print("开始加载知识库...")
        # 加载文档目录下的所有文档
        knowledge_base.load()
        print("知识库加载完成")
    except Exception as e:
        print(f"加载知识库时出错: {str(e)}")

if __name__ == "__main__":
    try:
        # 加载知识
        print("正在加载知识库...")
        load_knowledge()
        print("知识库加载完成!")
        
        # 测试查询
        while True:
            question = input("\n请输入您的问题(输入q退出): ")
            if question.lower() == 'q':
                break
            
            print("\n正在查询知识图谱...")
            answer = query_knowledge(question)
            print("\n回答:", answer)
            
    finally:
        # 关闭数据库连接
        print("\n正在关闭数据库连接...")
        graph_db.close()
        print("程序已结束!") 